from core.cuda import cuda_module
from core.config import Config
from core.function import Function
from utils.common import to_variable


# dropout 是一种通过随机删除（禁用）神经元进行训练的方法。模型在训练时随机选择隐藏层的神经元，并删除选中的神经元
class Dropout(Function):

    def dropout(self, x, dropout_ratio=0.5):
        """
        在训练阶段对输入变量应用 dropout。

        Args:
            x: 输入变量。
            dropout_ratio (float): Dropout 比率，将输入设置为零的概率。

        Returns:
            Variable: 应用 dropout 后的输出变量。

        """
        x = to_variable(x)

        if Config.train:
            xp = cuda_module
            mask = xp.random.rand(*x.shape) > dropout_ratio
            scale = xp.array(1.0 - dropout_ratio).astype(x.dtype)
            y = x * mask / scale
            return y
        else:
            return x
